#include "dicom/storage/encrypt_stream.h"

namespace dicom {

EncryptStream::EncryptStream(const std::string& hexkey, WriterStream* out)
    : cipher_(SDBasic::bytearray::ByteArray::fromHexStr(hexkey), 4096),
      src_size_(0),
      enc_size_(0),
      out_(out),
      hexkey_(hexkey) {
  zstrm_.zalloc = Z_NULL;
  zstrm_.zfree = Z_NULL;
  zstrm_.opaque = Z_NULL;
  ::deflateInit(&zstrm_, Z_DEFAULT_COMPRESSION);
}

EncryptStream::~EncryptStream() { ::deflateEnd(&zstrm_); }

int EncryptStream::Finalize() {
  int ret, flush;
  unsigned have;
  Status s;

  flush = Z_FINISH;
  zstrm_.avail_in = 0;
  zstrm_.next_in = NULL;

  do {
    zstrm_.avail_out = kBlockSize;
    zstrm_.next_out = outbuf_;
    ret = ::deflate(&zstrm_, flush);
    assert(ret != Z_STREAM_ERROR);
    have = kBlockSize - zstrm_.avail_out;
    switch (ret) {
      case Z_NEED_DICT:
      case Z_DATA_ERROR:
      case Z_MEM_ERROR:
        s = Status(Status::kOther, "zip error");
        goto error;
    }

    if (!cipher_.Update(outbuf_, have)) {
      s = Status(Status::kOther, "aes error");
      goto error;
    }

    if (cipher_.len() > 0) {
      int n = out_->Write(cipher_.data(), cipher_.len());
      if (n < 0 || n != cipher_.len()) {
        s = Status(Status::kOther, "failed to write");
        goto error;
      }
      enc_sha1_.Update(cipher_.data(), cipher_.len());
      enc_size_ += cipher_.len();
    }
  } while (zstrm_.avail_out == 0 && ret != Z_STREAM_END);

  if (!cipher_.Finalize()) {
    s = Status(Status::kOther, "AES256.Finalize() failed");
    goto error;
  }

  if (cipher_.len() > 0) {
    int n = out_->Write(cipher_.data(), cipher_.len());
    if (n < 0) {
      s = Status(Status::kOther, "failed to write");
      goto error;
    }

    enc_sha1_.Update(cipher_.data(), cipher_.len());
    enc_size_ += cipher_.len();
  }

  if (out_->Finalize() < 0) {
    s = Status(Status::kOther, "failed to finalize write");
    goto error;
  }

  src_md5_.Finalize();
  src_sha1_.Finalize();
  src_sha512_.Finalize();
  enc_sha1_.Finalize();
  return 0;

error:
  LOG_ERROR << s;
  return -1;
}

void EncryptStream::Close() { out_->Close(); }

void EncryptStream::GetHash(HashEntry* hash) {
  hash->src_md5 = src_md5_.digest();
  hash->src_size = src_size_;
  hash->src_digest = StringPrintf("%s%s%lx", src_md5_.digest().c_str(),
                                  src_sha1_.digest().c_str(), src_size_);
  hash->src_sha512 = src_sha512_.digest();
  hash->hexkey = hexkey_;
  hash->enc_size = enc_size_;
  hash->enc_digest = enc_sha1_.digest();
}

int EncryptStream::Write(const void* buf, size_t n) {
  int ret, flush;
  unsigned have;
  Status s;

  src_md5_.Update(buf, n);
  src_sha1_.Update(buf, n);
  src_sha512_.Update(buf, n);
  src_size_ += n;

  flush = Z_NO_FLUSH;
  zstrm_.avail_in = static_cast<uInt>(n);
  zstrm_.next_in = const_cast<Bytef*>(reinterpret_cast<const Bytef*>(buf));

  do {
    zstrm_.avail_out = kBlockSize;
    zstrm_.next_out = outbuf_;
    ret = ::deflate(&zstrm_, flush);
    assert(ret != Z_STREAM_ERROR);
    have = kBlockSize - zstrm_.avail_out;
    switch (ret) {
      case Z_NEED_DICT:
      case Z_DATA_ERROR:
      case Z_MEM_ERROR:
        s = Status(Status::kOther, "zip error");
        goto error;
    }

    if (!cipher_.Update(outbuf_, have)) {
      s = Status(Status::kOther, "aes error");
      goto error;
    }

    if (cipher_.len() > 0) {
      int n = out_->Write(cipher_.data(), cipher_.len());
      if (n < 0 || n != cipher_.len()) {
        s = Status(Status::kOther, "failed to write");
        goto error;
      }
      enc_sha1_.Update(cipher_.data(), cipher_.len());
      enc_size_ += cipher_.len();
    }
  } while (zstrm_.avail_out == 0);
  return static_cast<int>(n);

error:
  LOG_INFO << s;
  return -1;
}
}
